'''
Created on Aug 24, 2010

@author: alex
'''
import numpy as np 
from scipy.stats import beta, binom
from util import ranking


def sgnTestPValueH0( dL, twoTailed=True  ):
    dL = np.asarray(dL)
    
    kNeg = (dL<0).sum()
    kPos = (dL>0).sum()
    n = kNeg + kPos
    
    if twoTailed: k = min(kNeg, kPos)
    else : k = kPos 
    
    p = binom.cdf( k, n, 0.5 )
    
    if twoTailed :  return 2*p
    else :          return p


def wsrDistrH0(n):
    """
    Recursive formula to calculate the Wilcoxon signed rank distribution under H0 for n samples.
    """
    if n==0:
        return np.array([1])
    else:
        N = (n*(n+1))/2+1
        d = np.zeros( N )
        d_ = wsrDistrH0(n-1)/2.
        N_ = len(d_)
        d[:N_] = d_
        d[-N_:] += d_
        return d

def _wsrPValue( n, wPos, twoTailed=True ):
    """
    n : number of samples
    wPos : sum of positive rank
    twoTailed : whether or not we calculate the pValue of the two tailed
    """
    
    if twoTailed:
        wNeg = (n*(n+1))/2 - wPos
        w = min(wPos, wNeg)
    else:
        w = wPos

    distr = wsrDistrH0(n)
    p = distr[:w+1].sum()
    if twoTailed : return 2*p
    else : return p
    
def wsrTestPValueH0( dL, twoTailed=True ):
    """
    High level function to calculate the Wilcoxon signed rank p-values for
    a list of differences dL.
    towTailed : whether or not we calculate the pValue of the two tailed
    """
    
    dL= np.asarray( dL )
    rL = np.array(ranking( np.abs( dL ) ))

    wPos = rL[ dL > 0  ].sum()
    n = (dL != 0).sum()
    return _wsrPValue(n, wPos, twoTailed)



def uBinBound(m,k,delta=0.1, uniformPrior=False):
    return beta.ppf( 1-delta, k+1, m-k+ int(uniformPrior) )

def lBinBound(m,k,delta=0.1, uniformPrior=False):
    return beta.isf( 1-delta, k+int(uniformPrior) , m-k+1)
  


def binBound(m,k,delta=0.1, uniformPrior=False):
    
    assert abs(int(m) - m) < 1e-8
    assert abs(int(k) - k) < 1e-8
    assert delta <=0.5 and delta >0
    assert m >= k
    assert m > 0
    assert k >= 0

    return lBinBound(m,k,delta, uniformPrior), uBinBound(m,k,delta, uniformPrior) 
